package ca

import (
	"crypto"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"net/url"
	"os"
	"strings"
	"time"
	"gitee.com/shanyou/sshplus"
	"gitee.com/shanyou/sshplus/utils"
	"github.com/gravitational/trace"
	"golang.org/x/crypto/ssh"
)

const (
	CertUserRoleNormal = "normal"
	CertUserRoleAdmin  = "admin"
	CertUserRoleRobot  = "robot"

	// ca 证书类型，用户不同的证书认证区别于ssh.UserCert,在commons中通过UrlParamters的形式配置
	CertTypeUser = "user" //ssh.UserCert
	CertTypeHost = "host" //ssh.HostCert
)

const (
	// CertExtensionPermitX11Forwarding allows X11 forwarding for certificate
	CertExtensionPermitX11Forwarding = "permit-X11-forwarding"
	// CertExtensionPermitAgentForwarding allows agent forwarding for certificate
	CertExtensionPermitAgentForwarding = "permit-agent-forwarding"
	// CertExtensionPermitPTY allows user to request PTY
	CertExtensionPermitPTY = "permit-pty"
	// CertExtensionPermitPortForwarding allows user to request port forwarding
	CertExtensionPermitPortForwarding = "permit-port-forwarding"
	CertExtensionUserRole             = "user-role"
)

// HostCertParams defines all parameters needed to generate a host certificate
type HostCertParams struct {
	// HostID is used by Teleport to uniquely identify a node within a cluster
	HostID string
	// Principals is a list of additional principals to add to the certificate.
	Principals []string
	// NodeName is the DNS name of the node
	HostName string
	// TTL defines how long a certificate is valid for
	TTL       time.Duration
	Inception time.Time
}

// Check checks parameters for errors
func (c HostCertParams) Check() error {
	if c.TTL < time.Minute {
		return trace.BadParameter("TTL can't be less than %v", time.Minute)
	}

	if c.Inception.Add(c.TTL).Before(time.Now()) {
		return trace.BadParameter("inception %v too old", c.Inception)
	}
	if c.HostID == "" && len(c.Principals) == 0 {
		return trace.BadParameter("HostID [%q] or Principals [%q] are required",
			c.HostID, c.Principals)
	}

	return nil
}

// UserCertParams defines OpenSSH user certificate parameters
type UserCertParams struct {
	Inception time.Time
	// TTL defines how long a certificate is valid for
	TTL time.Duration
	// Username is teleport username
	Username string
	// AllowedLogins is a list of SSH principals
	AllowedLogins []string
	// PermitX11Forwarding permits X11 forwarding for this cert
	PermitX11Forwarding bool
	// PermitAgentForwarding permits agent forwarding for this cert
	PermitAgentForwarding bool
	// PermitPortForwarding permits port forwarding.
	PermitPortForwarding bool
	UserRole             string
}

func (c UserCertParams) Check() error {
	if c.TTL < time.Minute {
		return trace.BadParameter("TTL can't be less than %v", time.Minute)
	}

	if c.Inception.Add(c.TTL).Before(time.Now()) {
		return trace.BadParameter("inception %v too old", c.Inception)
	}
	if len(c.AllowedLogins) == 0 {
		return trace.BadParameter("AllowedLogins are required")
	}
	return nil
}

// CertAuthority ca的抽象，支持生成用户证书以及主机证书
type CertAuthority interface {
	//GenerateHostCert return serializes key for inclusion in an OpenSSH
	// authorized_keys file and ssh Signer for cert verify and sign
	GenerateHostCert(crypto.Signer, HostCertParams) ([]byte, ssh.Signer, error)
	GenerateUserCert(crypto.Signer, UserCertParams) ([]byte, ssh.Signer, error)
}

// CA 使用ecdsa进行证书生成与签名
type CA struct {
	signer crypto.Signer
}

func (c *CA) GenerateHostCert(key crypto.Signer, param HostCertParams) ([]byte, ssh.Signer, error) {
	if err := param.Check(); err != nil {
		return nil, nil, err
	}
	caPub, err := ssh.NewPublicKey(c.signer.Public())
	if err != nil {
		return nil, nil, trace.Wrap(err)
	}
	caSigner, err := ssh.NewSignerFromKey(c.signer)
	if err != nil {
		return nil, nil, trace.Wrap(err)
	}

	//如果没有提供签名用的公钥则提示错误
	if key == nil {
		return nil, nil, trace.Errorf("invalid input key")
	}

	publicKey, err := ssh.NewPublicKey(key.Public())
	if err != nil {
		return nil, nil, trace.Wrap(err)
	}
	keySigner, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return nil, nil, trace.Wrap(err)
	}

	// create principals
	principals := param.Principals
	if len(param.HostName) > 0 {
		principals = append(principals, param.HostName)
	}

	principals = utils.Deduplicate(principals)
	validAfter := uint64(time.Now().UTC().Add(-1 * time.Minute).Unix())
	if !param.Inception.IsZero() {
		validAfter = uint64(param.Inception.Unix())
	}
	validBefore := uint64(ssh.CertTimeInfinity)
	if param.TTL != 0 {
		b := time.Now().UTC().Add(param.TTL)
		validBefore = uint64(b.Unix())
	}

	// Create certificate and signer.
	cert := &ssh.Certificate{
		KeyId:           param.HostID,
		ValidPrincipals: principals,
		Key:             publicKey,
		SignatureKey:    caPub,
		ValidAfter:      validAfter,
		ValidBefore:     validBefore,
		CertType:        ssh.HostCert,
	}

	err = cert.SignCert(rand.Reader, caSigner)
	if err != nil {
		return nil, nil, trace.Wrap(err)
	}

	certSigner, err := ssh.NewCertSigner(cert, keySigner)
	return MarshalCertAuthorizedKey(CertTypeHost, cert), certSigner, err
}

func (c *CA) GenerateUserCert(key crypto.Signer, param UserCertParams) ([]byte, ssh.Signer, error) {
	if err := param.Check(); err != nil {
		return nil, nil, err
	}
	caPub, err := ssh.NewPublicKey(c.signer.Public())
	if err != nil {
		return nil, nil, trace.Wrap(err)
	}
	caSigner, err := ssh.NewSignerFromKey(c.signer)
	if err != nil {
		return nil, nil, trace.Wrap(err)
	}

	//如果没有提供签名用的公钥则提示错误
	if key == nil {
		return nil, nil, trace.Errorf("invalid input key")
	}

	publicKey, err := ssh.NewPublicKey(key.Public())
	if err != nil {
		return nil, nil, trace.Wrap(err)
	}
	keySigner, err := ssh.NewSignerFromKey(key)
	if err != nil {
		return nil, nil, trace.Wrap(err)
	}

	validAfter := uint64(time.Now().UTC().Add(-1 * time.Minute).Unix())
	if !param.Inception.IsZero() {
		validAfter = uint64(param.Inception.Unix())
	}

	validBefore := uint64(ssh.CertTimeInfinity)
	if param.TTL != 0 {
		b := time.Now().UTC().Add(param.TTL)
		validBefore = uint64(b.Unix())
	}
	allowLogins := utils.Deduplicate(param.AllowedLogins)
	// Create certificate and signer.
	cert := &ssh.Certificate{
		KeyId:           param.Username,
		ValidPrincipals: allowLogins,
		Key:             publicKey,
		SignatureKey:    caPub,
		ValidAfter:      validAfter,
		ValidBefore:     validBefore,
		CertType:        ssh.UserCert,
	}

	cert.Permissions.Extensions = map[string]string{
		CertExtensionPermitPTY: "",
		// CertExtensionPermitPortForwarding: "",
	}
	if param.PermitX11Forwarding {
		cert.Permissions.Extensions[CertExtensionPermitX11Forwarding] = ""
	}
	if param.PermitAgentForwarding {
		cert.Permissions.Extensions[CertExtensionPermitAgentForwarding] = ""
	}
	if param.PermitPortForwarding {
		cert.Permissions.Extensions[CertExtensionPermitPortForwarding] = ""
	}
	if len(param.UserRole) > 0 {
		cert.Permissions.Extensions[CertExtensionUserRole] = param.UserRole
	}

	err = cert.SignCert(rand.Reader, caSigner)
	if err != nil {
		return nil, nil, trace.Wrap(err)
	}

	certSigner, err := ssh.NewCertSigner(cert, keySigner)
	return MarshalCertAuthorizedKey(CertTypeUser, cert), certSigner, err
}

// NewCA create new ca
func NewCA(signer crypto.Signer) (*CA, error) {
	return &CA{
		signer: signer,
	}, nil
}

// LoadCA 通过文件加载CA的PrivateKey
func LoadCAFromFile(path string) (*CA, error) {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return nil, trace.Wrap(err)
	}
	privPem, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}

	return LoadCAFromBytes(privPem)
}

func LoadCAFromBytes(keybytes []byte) (*CA, error) {
	block, _ := pem.Decode(keybytes)
	if x509.IsEncryptedPEMBlock(block) {
		return nil, trace.Errorf("Failed decryptPEM PEM. it's not a decryped PEM")
	}
	privateKey, err := x509.ParseECPrivateKey(block.Bytes)

	if err != nil {
		return nil, trace.Wrap(err)
	}

	ca := &CA{
		signer: privateKey,
	}
	return ca, nil
}

// MarshalCertAuthorizedKey 新增comment 协议
func MarshalCertAuthorizedKey(certType string, key ssh.PublicKey) []byte {
	comments := url.Values{
		"type":     []string{certType},
		"protocal": []string{sshplus.ProtocolVersion},
	}
	keyBytes := ssh.MarshalAuthorizedKey(key)
	res := fmt.Sprintf("%s %s", strings.TrimSpace(string(keyBytes)), comments.Encode())
	return []byte(res)
}
